'''
instance跟踪算法
author: xiaohe
time: 2022/05/31
'''
import numpy as np
# from sort.roi_feature_extractor import feature_extracor_network
from simple_extractor import simple_extractor as feature_extracor_network

class MyTracker(object):
    def __init__(self, score_threshold=0.5, max_id=200, max_age=5, momentum=0.8, **kwargs) -> None:
        self.max_id = max_id
        self.max_age = max_age
        self.score_th = score_threshold
        self.model = feature_extracor_network(**kwargs)
        
        # 历史有效的特征
        self.exist_feature_vetors = None
        # 追踪框生命计数, <=0 时将被移除
        self.tracker_life = []
        # 每一个框的追踪id
        self.trace_ids = []
        # 变量更新momentum
        self.momentum = momentum
        # 追踪计数变量
        self.cur_trace_id = 1
        
    def __call__(self, bboxes, scores, class_ids, image):
        # t0 = time.time()
        # 抽特征
        features = self.get_features(image, bboxes)
        if self.exist_feature_vetors is None:
            # 初始化
            self.exist_feature_vetors = features
            self.tracker_life = np.array([self.max_age]*len(features))
            self.trace_ids = np.array([i for i in range(len(features))])
            self.cur_trace_id = len(self.trace_ids)
        
        # t1 = time.time()
        # 计算相似度分数
        similarity_scores = self.compute_similarity(features, self.exist_feature_vetors)

        # t2 = time.time()
        # 进行配对
        pair_list = self.get_match_pair(similarity_scores)

        # t3 = time.time()
        # 更新tracker状态
        _pair_list = self.update_tracker_state(pair_list, features)

        # t4 = time.time()
        # 分配追踪id号
        instance_ids = [-1]*len(bboxes)
        instance_ids = self.assign_trace_id(_pair_list, instance_ids)

        outputs = np.hstack([bboxes, instance_ids[:, np.newaxis], 
                            class_ids[:, np.newaxis], scores[:, np.newaxis]])
        # t5 = time.time()
        
        # print(t5-t0)
        # print("{:^4.2f}, {:^4.2f}, {:^4.2f}, {:^4.2f}, {:^4.2f}".format(t1-t0, t2-t1, t3-t2, t4-t3, t5-t4))
        return outputs

    def assign_trace_id(self, match_pair, instance_ids):
        for i in range(len(match_pair)):
            if match_pair[i] == -1:
                continue
            instance_ids[match_pair[i]] = self.trace_ids[i]
        return np.array(instance_ids, dtype=np.int32)

    def update_tracker_state(self, match_pair, features):
        match_pair = np.array(match_pair, dtype=np.int32)
        # 修正tracker生命周期
        active_indexes = []
        for i in range(len(match_pair)):
            if match_pair[i] == -1:
                # 如果该历史feature此次没有匹配到框
                self.tracker_life[i] -= 1
                if self.tracker_life[i] > 0:
                    active_indexes.append(i)
            else:
                # 如果该历史feature此次匹配到了框
                self.tracker_life[i] = self.max_age
                active_indexes.append(i)
                # 修正feature, 动量更新
                self.exist_feature_vetors[i] = self.momentum*self.exist_feature_vetors[i] + (1-self.momentum)*features[match_pair[i]]

        # 去除不活跃的历史tracker
        self.exist_feature_vetors = self.exist_feature_vetors[active_indexes]
        self.tracker_life = self.tracker_life[active_indexes]
        self.trace_ids = self.trace_ids[active_indexes]
        match_pair = match_pair[active_indexes]
        
        # 新增加的
        for i in range(len(features)):
            if i in match_pair:
                # 如果已经匹配到则跳过
                continue
            match_pair = np.hstack([match_pair, i])
            self.exist_feature_vetors = np.vstack([self.exist_feature_vetors, [features[i]]])
            self.tracker_life = np.hstack([self.tracker_life, self.max_age])
            self.trace_ids = np.hstack([self.trace_ids, self.cur_trace_id])
            self.cur_trace_id = (self.cur_trace_id + 1)%self.max_id

        return match_pair

    def get_match_pair(self, similarity_scores_matrix):
        '''
        匈牙利算法配对
        :param similarity_scores_matrix: (N1, N2), N1为检测框数量,N2为已经存在的检测框数量
        :return map_list, 
        '''
        # 本轮检测到的(A), 已经检测到的(B)
        detect_num, exist_num = similarity_scores_matrix.shape[0], similarity_scores_matrix.shape[1]
        
        # B被分配到哪一个A, -1代表没有分配对象
        b2a = [-1]*exist_num

        # B元素是否已经被访问过
        b_visited = [False]*exist_num

        sort_keys = [i for i in range(exist_num)]
        def __match(a_index:int):
            scores = similarity_scores_matrix[a_index]
            key = sorted(sort_keys, key=lambda x:scores[x], reverse=True)
            # for b_index in range(exist_num):
            for b_index in key:
                if(scores[b_index] > self.score_th and (not b_visited[b_index])):
                    b_visited[b_index] = True
                    if(b2a[b_index] == -1 or 
                        (scores[b_index] > np.max(similarity_scores_matrix[b2a[b_index]]) and __match(b2a[b_index]))):
                    # if(b2a[b_index] == -1 or __match(b2a[b_index])):
                    # if(b2a[b_index] == -1):
                        b2a[b_index] = a_index
                        return True
            return False

        for i in range(detect_num):
            b_visited = [False]*exist_num
            __match(i)

        return b2a
    
    def compute_similarity(self, A, B):
        '''
        计算余弦相似度
        :param Matrix A (N1, M)
        :param Matrix B (N2, M)
        :return Matrix Y (N1, N2)
        '''
        Y = np.dot(A, B.T) \
            / np.dot(np.linalg.norm(A, axis=1, keepdims=True), 
                      np.linalg.norm(B, axis=1, keepdims=True).T)
        return Y

    def get_features(self, ori_img, bboxes):
        '''
        提取区域特征
        '''
        im_crops = []
        for box in bboxes:
            x1, y1, x2, y2 = [int(i) for i in box]
            im = ori_img[y1:y2, x1:x2]
            im_crops.append(im)
        if im_crops:
            features = self.model(im_crops)
        else:
            features = np.array([])
        return features